"""
视觉处理模块 - 负责捕获和分析魔兽争霸3游戏画面

该模块使用OpenCV库进行图像处理，识别游戏中的关键元素，
如资源数量、单位状态、建筑位置等，为AI决策提供视觉输入。
"""

import cv2
import numpy as np
import pyautogui
import time
from typing import Tuple, Optional, Dict, Any


class VisionProcessor:
    """
    视觉处理类 - 处理魔兽争霸3游戏画面
    
    该类负责：
    1. 捕获游戏画面
    2. 识别游戏中的关键元素
    3. 提取游戏状态信息
    """
    
    def __init__(self, screen_region: Tuple[int, int, int, int] = (0, 0, 1920, 1080)):
        """
        初始化视觉处理器
        
        Args:
            screen_region: 屏幕捕获区域 (x, y, width, height)
        """
        # 屏幕捕获区域
        self.screen_region = screen_region
        
        # 游戏元素模板（用于模板匹配）
        self.templates = {}
        
        # 颜色阈值（用于资源识别）
        self.color_thresholds = {
            'gold': ([0, 100, 100], [20, 255, 255]),  # 金色资源
            'wood': ([35, 50, 50], [85, 255, 255]),   # 木材资源
        }
        
        print("视觉处理器初始化完成")
    
    def capture_screen(self) -> Optional[np.ndarray]:
        """
        捕获屏幕指定区域的图像
        
        Returns:
            np.ndarray: 捕获的屏幕图像，如果失败返回None
        """
        try:
            # 使用pyautogui捕获屏幕
            screenshot = pyautogui.screenshot(region=self.screen_region)
            # 转换为OpenCV格式
            frame = cv2.cvtColor(np.array(screenshot), cv2.COLOR_RGB2BGR)
            return frame
        except Exception as e:
            print(f"屏幕捕获失败: {e}")
            return None
    
    def load_template(self, name: str, template_path: Optional[str] = None) -> bool:
        """
        加载图像模板用于匹配
        
        Args:
            name: 模板名称
            template_path: 模板文件路径（可选）
            
        Returns:
            bool: 加载成功返回True，否则返回False
        """
        try:
            # 如果未提供路径，友好地返回失败而不抛出异常，兼容测试单参调用
            if not template_path:
                print(f"未提供模板路径，无法加载模板: {name}")
                return False

            template = cv2.imread(template_path, cv2.IMREAD_COLOR)
            if template is not None:
                self.templates[name] = template
                return True
            else:
                print(f"无法加载模板: {template_path}")
                return False
        except Exception as e:
            print(f"加载模板失败 {template_path}: {e}")
            return False
    
    def match_template(self, frame: np.ndarray, template_name: str, 
                      threshold: float = 0.8) -> Optional[Tuple[int, int, int, int]]:
        """
        在图像中匹配指定模板
        
        Args:
            frame: 要搜索的图像
            template_name: 模板名称
            threshold: 匹配阈值
            
        Returns:
            Tuple[int, int, int, int]: 匹配区域的坐标(x, y, w, h)，未找到返回None
        """
        if template_name not in self.templates:
            print(f"模板 {template_name} 未加载")
            return None
        
        template = self.templates[template_name]
        
        try:
            # 执行模板匹配
            result = cv2.matchTemplate(frame, template, cv2.TM_CCOEFF_NORMED)
            # 查找最佳匹配位置
            min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
            
            if max_val >= threshold:
                # 返回匹配区域坐标 (x, y, w, h)
                return (max_loc[0], max_loc[1], template.shape[1], template.shape[0])
            else:
                return None
        except Exception as e:
            print(f"模板匹配失败: {e}")
            return None
    
    def detect_resources(self, frame: np.ndarray) -> Dict[str, int]:
        """
        检测游戏资源数量（金币、木材）
        
        Args:
            frame: 游戏画面帧
            
        Returns:
            Dict[str, int]: 资源数量字典
        """
        resources = {'gold': 0, 'wood': 0}
        
        try:
            # 转换为HSV颜色空间便于颜色检测
            hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
            
            # 检测金币
            gold_mask = cv2.inRange(hsv, 
                                   np.array(self.color_thresholds['gold'][0]), 
                                   np.array(self.color_thresholds['gold'][1]))
            resources['gold'] = np.sum(gold_mask > 0)
            
            # 检测木材
            wood_mask = cv2.inRange(hsv, 
                                   np.array(self.color_thresholds['wood'][0]), 
                                   np.array(self.color_thresholds['wood'][1]))
            resources['wood'] = np.sum(wood_mask > 0)
            
        except Exception as e:
            print(f"资源检测失败: {e}")
        
        return resources
    
    def detect_minimap(self, frame: np.ndarray) -> Optional[np.ndarray]:
        """
        提取小地图区域
        
        Args:
            frame: 游戏画面帧
            
        Returns:
            np.ndarray: 小地图图像，如果未找到返回None
        """
        # 小地图通常在右上角，我们假设它在固定位置
        # 这里需要根据实际游戏画面调整坐标
        minimap_region = (frame.shape[1] - 200, 0, 200, 200)  # 右上角200x200区域
        
        try:
            # 提取小地图区域
            minimap = frame[minimap_region[1]:minimap_region[1]+minimap_region[3],
                           minimap_region[0]:minimap_region[0]+minimap_region[2]]
            return minimap
        except Exception as e:
            print(f"小地图提取失败: {e}")
            return None
    
    def analyze_game_state(self) -> Dict[str, Any]:
        """
        分析当前游戏状态
        
        Returns:
            Dict[str, Any]: 游戏状态信息
        """
        state = {
            'timestamp': time.time(),
            'resources': {},
            'minimap': None,
            'units': [],
            'buildings': []
        }
        
        # 捕获当前画面
        frame = self.capture_screen()
        if frame is None:
            print("无法捕获游戏画面")
            return state
        
        # 检测资源
        state['resources'] = self.detect_resources(frame)
        
        # 提取小地图
        state['minimap'] = self.detect_minimap(frame)
        
        # 这里可以添加更多分析逻辑
        # 例如：检测单位、建筑、敌人位置等
        
        return state


# 使用示例
if __name__ == "__main__":
    # 创建视觉处理器实例
    vision = VisionProcessor()
    
    # 分析游戏状态
    game_state = vision.analyze_game_state()
    print("游戏状态:", game_state)
